package cn.iocoder.yudao.framework.mybatis.component;

import com.zaxxer.hikari.HikariDataSource;
import org.apache.ibatis.cache.CacheKey;
import org.apache.ibatis.executor.Executor;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.mapping.ParameterMapping;
import org.apache.ibatis.mapping.ParameterMode;
import org.apache.ibatis.plugin.Interceptor;
import org.apache.ibatis.plugin.Intercepts;
import org.apache.ibatis.plugin.Invocation;
import org.apache.ibatis.plugin.Signature;
import org.apache.ibatis.reflection.MetaObject;
import org.apache.ibatis.session.Configuration;
import org.apache.ibatis.session.ResultHandler;
import org.apache.ibatis.session.RowBounds;
import org.apache.ibatis.type.TypeHandlerRegistry;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.sql.DataSource;
import java.lang.reflect.Method;
import java.sql.Timestamp;
import java.time.LocalDate;
import java.time.LocalDateTime;
import java.util.ArrayList;
import java.util.Date;
import java.util.List;
import java.util.StringTokenizer;

/**
 * TODO
 *
 * @className: SqlCostInterceptor
 * @author: tangkh
 * @date: 2023-09-19 16:52
 */
@Intercepts({@Signature(
        type = Executor.class,
        method = "update",
        args = {MappedStatement.class, Object.class}
), @Signature(
        type = Executor.class,
        method = "query",
        args = {MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class, CacheKey.class, BoundSql.class}
), @Signature(
        type = Executor.class,
        method = "query",
        args = {MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class}
)})
public class SqlCostInterceptor implements Interceptor {
    private static final Logger log = LoggerFactory.getLogger(SqlCostInterceptor.class);
    public static final String DS_ITEM_DATA_SOURCE = "com.baomidou.dynamic.datasource.ds.ItemDataSource";
    public static final String DATA_SOURCE_CONTEXT_HOLDER = "com.baomidou.dynamic.datasource.toolkit.DynamicDataSourceContextHolder";
    public static final String DYNAMIC_ROUTING_DATA_SOURCE_CLASS = "cn.iocoder.yudao.framework.datasource.config.DynamicRoutingDataSource";
    public static final String DATA_SOURCE = "getDataSource";
    public static final String PEEK = "peek";
    private static boolean dynamicRoutingDataSourceFlag = true;
    private static Class<?> cls = null;

    public SqlCostInterceptor() {
    }

    @Override
    public Object intercept(Invocation invocation) throws Throwable {
        long startTime = System.currentTimeMillis();
        BoundSql boundSql = null;
        Object[] args = null;
        String jdbcUrl = "";

        Object var21;
        try {
            args = invocation.getArgs();
            MappedStatement ms = (MappedStatement) args[0];
            Object parameter = null;
            if (invocation.getArgs().length > 1) {
                parameter = invocation.getArgs()[1];
            }

            boundSql = ms.getBoundSql(parameter);
            DataSource dataSource = ms.getConfiguration().getEnvironment().getDataSource();
            if (dataSource instanceof HikariDataSource) {
                jdbcUrl = ((HikariDataSource) dataSource).getJdbcUrl();
            } else if (dynamicRoutingDataSourceFlag && DYNAMIC_ROUTING_DATA_SOURCE_CLASS.equals(dataSource.getClass().getCanonicalName())) {
                try {
                    Class<?> clsHolder = Class.forName(DATA_SOURCE_CONTEXT_HOLDER);
                    Method methodHolder = clsHolder.getDeclaredMethod("peek");
                    String peek = (String) methodHolder.invoke(methodHolder, (Object[]) null);
                    Method method = cls.getDeclaredMethod(DATA_SOURCE, String.class);
                    DataSource ds = (DataSource) method.invoke(dataSource, peek);
                    if (DS_ITEM_DATA_SOURCE.equals(ds.getClass().getCanonicalName())) {
                        Method methodHsource1 = ds.getClass().getDeclaredMethod(DATA_SOURCE, (Class[]) null);
                        jdbcUrl = ((HikariDataSource) methodHsource1.invoke(ds, (Object[]) null)).getJdbcUrl();
                    }
                } catch (Exception var19) {
                    log.warn("jdbc url解析错误", var19.getMessage());
                }
            }

            var21 = invocation.proceed();
        } finally {
            this.printSqlLog(args, startTime, boundSql, jdbcUrl);
        }

        return var21;
    }

    public void printSqlLog(Object[] args, long startTime, BoundSql boundSql, String jdbUrl) {
        if (boundSql != null) {
            long endTime = System.currentTimeMillis();
            long sqlCost = endTime - startTime;
            Object parameterObject = boundSql.getParameterObject();
            MappedStatement ms = (MappedStatement) args[0];
            if (args.length == 6) {
                boundSql = (BoundSql) args[5];
            }

            String sql = boundSql.getSql();
            Configuration configuration = ms.getConfiguration();
            String allSql = null;

            try {
                allSql = this.getFllSql(configuration, boundSql, parameterObject);
            } catch (Exception var16) {
                log.warn("sql 格式化失败:{}", var16.getMessage());
            }

            if (allSql == null || allSql.length() == 0) {
                allSql = sql;
            }

            String dbName = this.jdbcUrl2DbName(jdbUrl);
            if (sqlCost > 200L) {
                log.warn("--SQL耗时{}ms.数据库名：{}，连接详情:{},\r\n{}", sqlCost, dbName, jdbUrl, allSql);
            } else {
                log.info("--SQL耗时{}ms,数据库名：{}，连接详情:{},\r\n{}", sqlCost, dbName, jdbUrl, allSql);
            }

        }
    }

    private String jdbcUrl2DbName(String jdbUrl) {
        if (!jdbUrl.contains("/")) {
            return "";
        } else {
            int end = jdbUrl.length();
            if (jdbUrl.contains("?")) {
                end = jdbUrl.indexOf("?");
            }

            String dbName = jdbUrl.substring(jdbUrl.lastIndexOf("/") + 1, end);
            return dbName;
        }
    }

    public String getFllSql(Configuration configuration, BoundSql boundSql, Object parameterObject) {
        TypeHandlerRegistry typeHandlerRegistry = configuration.getTypeHandlerRegistry();
        List<ParameterMapping> parameterMappings = boundSql.getParameterMappings();
        if (parameterMappings != null && !parameterMappings.isEmpty()) {
            List<Object> paramList = new ArrayList<>();

            for (int i = 0; i < parameterMappings.size(); ++i) {
                ParameterMapping parameterMapping = parameterMappings.get(i);
                if (parameterMapping.getMode() != ParameterMode.OUT) {
                    Object value = null;
                    String propertyName = parameterMapping.getProperty();
                    if (boundSql.hasAdditionalParameter(propertyName)) {
                        value = boundSql.getAdditionalParameter(propertyName);
                    } else if (parameterObject == null) {
                        value = null;
                    } else if (typeHandlerRegistry.hasTypeHandler(parameterObject.getClass())) {
                        value = parameterObject;
                    } else {
                        MetaObject metaObject = configuration.newMetaObject(parameterObject);
                        value = metaObject.getValue(propertyName);
                    }

                    if (value instanceof Date) {
                        Timestamp timestamp = new Timestamp(((Date) value).getTime());
                        paramList.add(timestamp);
                    } else {
                        paramList.add(value);
                    }
                }
            }

            String sql = this.removeBreakingWhitespace(boundSql.getSql());
            StringTokenizer sqlTokenizer = new StringTokenizer(sql, "?");
            StringBuilder builder = new StringBuilder();
            int i = 0;

            while (sqlTokenizer.hasMoreTokens()) {
                builder.append(sqlTokenizer.nextToken());
                if (i < paramList.size()) {
                    builder.append(this.formatParamValue(paramList.get(i)));
                    ++i;
                }
            }

            if (i > 0) {
                sql = builder.toString();
            }

            return sql;
        } else {
            return null;
        }
    }

    public String formatParamValue(Object paramValue) {
        if (paramValue == null) {
            return "null";
        } else {
            if (paramValue instanceof String || paramValue instanceof Timestamp || paramValue instanceof LocalDate || paramValue instanceof LocalDateTime) {
                paramValue = "'" + paramValue + "'";
            }

            return paramValue.toString();
        }
    }

    protected String removeBreakingWhitespace(String original) {
        StringTokenizer whitespaceStripper = new StringTokenizer(original);
        StringBuilder builder = new StringBuilder();

        while (whitespaceStripper.hasMoreTokens()) {
            builder.append(whitespaceStripper.nextToken());
            builder.append(" ");
        }

        return builder.toString();
    }

    static {
        try {
            cls = Class.forName(DYNAMIC_ROUTING_DATA_SOURCE_CLASS);
        } catch (ClassNotFoundException var1) {
            dynamicRoutingDataSourceFlag = false;
        }

    }
}
